# coding:utf-8
# writingtime:2022-7-4
# author:wanjun

import os.path
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator

colors = ['brown', 'orange', 'purple', 'olive', 'gray', 'pink', 'black', 'red', 'green', 'blue', 'cyan', 'm']
markers = ['o', 's', 'd', 'p', 'h', 'x', 'v', '<', '.', '1', '*', '+']


def matrixReverse(matrix):
    '''
    function: 功能函数
    :param matrix: 矩阵
    :return: 行列交换的矩阵
    '''
    newMatrix = []
    for x in range(len(matrix[0])):
        temp = []
        for y in range(len(matrix)):
            temp.append(matrix[y][x])
        newMatrix.append(temp)

    return newMatrix

class Plot3D:

    def heatMap(self,x_list,y_list,z_matrix,title='title',xlabel='x',ylabel='y',zlabel='value',filename='text.png',
                imgShowing=True,imgSaving=False):
        '''
        function: 绘制三维热力图
        :param x_list: 一般为q的范围
        :param y_list: 该算子的其他参数范围
        :param z_matrix: 数值矩阵
        :param title: 标题
        :param xlabel: x轴的标签
        :param ylabel: y轴的标签
        :param filename: 文件名
        :param imgShowing: 是否显示图片
        :param imgSaving: 是否保存图片
        :return: null
        '''
        fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(8, 5))
        # 设置标题
        plt.title(title)
        # 设置x轴y轴z轴标签
        ax.set(xlabel=xlabel, ylabel=ylabel, zlabel=zlabel)
        # 算子其他参数的个数
        y_length=len(z_matrix)
        # q的个数
        x_length=len(x_list)
        # 将x，y，z的维度转换一致
        data_X = np.array([np.array(x_list) for i in range(y_length)])
        data_Y = np.array([np.array([y_list[j] for _ in range(x_length)]) for j in range(y_length)])
        data_Z = np.array([np.array(z_matrix[i]) for i in range(y_length)])
        # 画图
        surf = ax.plot_surface(data_X, data_Y, data_Z, cmap=plt.get_cmap('rainbow'), linewidth=0, antialiased=True)
        # 设置z轴的范围
        ax.set_zlim(-1.01, 1.01)
        # 将z轴范围等分
        ax.zaxis.set_major_locator(LinearLocator(10))
        # 保留1位小数
        ax.zaxis.set_major_formatter('{x:.1f}')
        # 设置x轴与y轴的坐标
        ax.xaxis.set_major_locator(mpl.ticker.FixedLocator(x_list))
        ax.yaxis.set_major_locator(mpl.ticker.FixedLocator(y_list))
        # 色条标签
        fig.colorbar(surf, shrink=0.5, aspect=5)
        if imgShowing:
            plt.show()
        if imgSaving:
            plt.savefig(filename)

if __name__=='__main__':
    a=Plot3D()
    x = [2, 3, 4, 5, 6]
    y = [1,2,3,4]
    z = [   [-0.02386906396874998, 0.755770521112178, -0.053917654786754005, -0.0036360276480000148, -0.39291258544128],
            [1, 1, 1, 1, 1],
            [-0.471508586764984, -0.13494840814902603, -0.259925952766329, 0.382493613993634, 0.5293814338602819],
            [-0.08359642452787201, -0.08099105757195202, 0.36077961664049807, -0.118141841729106, -0.04767854593041401]]
    a.heatMap(x,y,z)
